import os
import pandas as pd
import baostock as bs
from datetime import datetime
import time # 导入time模块用于延时

def get_market_data(cache_dir, symbol='sh.000300', adjust_flag="1"):
    """获取大盘指数数据，并进行缓存。"""
    fname = os.path.join(cache_dir, f"{symbol}_market_data.feather")
    if os.path.exists(fname):
        return pd.read_feather(fname)
    
    print(f"缓存未找到，正在从Baostock获取大盘指数 {symbol} 的数据...")
    today = datetime.now().strftime('%Y-%m-%d')
    
    # --- 数据获取逻辑直接在此处实现 ---
    fields = "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST"
    try:
        rs = bs.query_history_k_data_plus(
            symbol, fields,
            start_date="2015-01-01", end_date=today,
            frequency="d", adjustflag=adjust_flag
        )
        if rs.error_code != '0':
            print(f"警告: Baostock API在获取 {symbol} 数据时返回错误: {rs.error_msg}")
            return None
        data_list = [rs.get_row_data() for _ in iter(lambda: rs.next(), False)]
        if not data_list: return pd.DataFrame()
        
        df = pd.DataFrame(data_list, columns=rs.fields)
        numeric_cols = ['open', 'high', 'low', 'close', 'preclose', 'volume', 'amount', 'turn', 'pctChg']
        for col in numeric_cols:
            df[col] = pd.to_numeric(df[col], errors='coerce')
        df['date'] = pd.to_datetime(df['date'])
        df = df[df['tradestatus'] == '1'].copy()
        df.drop(columns=['tradestatus'], inplace=True)
        # --- 获取逻辑结束 ---

        if df is not None and not df.empty:
            df.to_feather(fname)
            print("大盘数据已缓存。")
        return df
    except Exception as e:
        print(f"错误: 在 get_market_data 中获取 {symbol} 数据时发生异常: {e}")
        return None


def prepare_stock_data(stock_code, start_date, end_date, cache_dir, adjust_flag="1"):
    """
    准备单只股票的数据，如果缓存不存在，则从Baostock获取并保存。
    """
    # 修复：如果代码已经包含 'sh' 或 'sz' 前缀，则直接使用；否则，根据规则添加前缀。
    if 'sh.' in stock_code or 'sz.' in stock_code:
        symbol = stock_code
    else:
        symbol = f"sh.{stock_code}" if stock_code.startswith('6') else f"sz.{stock_code}"
    
    fpath = os.path.join(cache_dir, f"{symbol}.feather")

    if not os.path.exists(fpath):
        # --- 数据获取逻辑直接在此处实现 ---
        fields = "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST"
        
        # --- 新增：带延时的重试逻辑 ---
        for i in range(3): # 最多重试3次
            try:
                rs = bs.query_history_k_data_plus(
                    symbol, fields,
                    start_date=start_date, end_date=end_date,
                    frequency="d", adjustflag=adjust_flag
                )
                if rs.error_code == '0':
                    # 成功获取数据，跳出重试循环
                    break
                else:
                    # BaoStock返回错误，记录并准备重试
                    print(f"警告: 第{i+1}次获取 {symbol} 数据失败, Baostock返回: {rs.error_msg}")
            except Exception as e:
                # 捕获网络连接等异常
                print(f"警告: 第{i+1}次获取 {symbol} 时发生网络异常: {e}")
                rs = None # 确保在异常后rs为None
            
            # 如果不是最后一次尝试，则延时后重试
            if i < 2:
                time.sleep(2) # 修复：增加重试延时到2秒
        # --- 重试逻辑结束 ---

        # 检查最终结果
        if rs is None or rs.error_code != '0':
            print(f"错误: 重试3次后，获取 {symbol} 数据最终失败。")
            return

        try:
            data_list = [rs.get_row_data() for _ in iter(lambda: rs.next(), False)]
            if not data_list: 
                # 如果没有获取到任何数据，也返回
                return

            stock_df = pd.DataFrame(data_list, columns=rs.fields)
            numeric_cols = ['open', 'high', 'low', 'close', 'preclose', 'volume', 'amount', 'turn', 'pctChg']
            for col in numeric_cols:
                stock_df[col] = pd.to_numeric(stock_df[col], errors='coerce')
            stock_df['date'] = pd.to_datetime(stock_df['date'])
            stock_df = stock_df[stock_df['tradestatus'] == '1'].copy()
            stock_df.drop(columns=['tradestatus'], inplace=True)
            # --- 获取逻辑结束 ---

            if stock_df is not None and not stock_df.empty:
                stock_df.to_feather(fpath)
        except Exception as e:
            # 捕获其他异常并打印，有助于调试
            print(f"错误: 在处理 {symbol} 的返回数据时发生异常: {e}")
            pass


def get_stock_list(source='zz500'):
    """
    获取股票列表。目前支持'zz500'（中证500）和'hs300'（沪深300）。
    :param source: 'zz500' or 'hs300'
    :return: 股票代码列表
    """
    print(f"正在从Baostock获取 {source} 成分股列表...")
    
    if source == 'zz500':
        rs = bs.query_zz500_stocks()
    elif source == 'hs300':
        rs = bs.query_hs300_stocks()
    else:
        print(f"警告: 不支持的股票列表源 '{source}'。")
        return []
        
    if rs.error_code != '0':
        print(f"警告: Baostock API在获取 {source} 成分股时返回错误: {rs.error_msg}")
        return []

    stock_list = [row_data[1] for row_data in rs.data]
    print(f"成功获取 {len(stock_list)} 只 {source} 成分股。")
    return stock_list


def filter_sub_new_stocks(stock_codes, required_start_date):
    """
    根据上市日期（IPO）过滤掉次新股。采用批量获取方式进行性能优化。
    
    :param stock_codes: 初始的股票代码列表 (e.g., ['sh.600000', ...])
    :param required_start_date: 回测所需的最早数据日期 (datetime.date object)
    :return: 过滤后的股票代码列表
    """
    print(f"\n--- 正在进行次新股过滤，要求上市日期早于 {required_start_date} ---")
    
    # 1. 一次性获取所有股票的基本信息
    rs = bs.query_stock_basic()
    if rs.error_code != '0':
        print(f"警告: Baostock API在获取所有股票基本信息时返回错误: {rs.error_msg}")
        # 在API失败时，保守地返回原始列表，避免过滤掉所有股票
        return stock_codes
        
    data_list = [rs.get_row_data() for _ in iter(lambda: rs.next(), False)]
    if not data_list:
        print("警告: 未能从Baostock获取到任何股票基本信息。")
        return stock_codes
        
    all_stocks_df = pd.DataFrame(data_list, columns=rs.fields)
    
    # 2. 数据清洗和类型转换
    all_stocks_df = all_stocks_df[all_stocks_df['type'] == '1'] # 只保留股票
    all_stocks_df['ipoDate'] = pd.to_datetime(all_stocks_df['ipoDate'], errors='coerce').dt.date
    all_stocks_df.dropna(subset=['ipoDate'], inplace=True)

    # 3. 将输入列表转换为DataFrame，以便合并
    target_stocks_df = pd.DataFrame(stock_codes, columns=['code'])

    # 4. 合并与过滤
    # 使用 inner join，只保留在两个DataFrame中都存在的股票
    merged_df = pd.merge(target_stocks_df, all_stocks_df[['code', 'ipoDate']], on='code', how='inner')
    
    # 执行日期过滤
    filtered_df = merged_df[merged_df['ipoDate'] < required_start_date]
    
    filtered_stocks = filtered_df['code'].tolist()
    
    print(f"次新股过滤完成，剩余 {len(filtered_stocks)} 只股票。")
    return filtered_stocks